package main

import (
	_ "embed"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
)

type Config struct {
	Port         int           `json:"port"`
	ReadTimeout  time.Duration `json:"read-timeout"`
	WriteTimeout time.Duration `json:"write-timeout"`
	Routes       []string      `json:"routes"`
}

//go:embed banner.txt
var banner string
var Cfg Config

func init() {
	log.SetPrefix("[yl-proxy] ")
	bytes, err := os.ReadFile("config.json")
	if err != nil {
		log.Fatalf("Error reading config.json: %s", err)
	}
	err = json.Unmarshal(bytes, &Cfg)
	if err != nil {
		log.Fatalf("Error parsing config.json: %s", err)
	}
}

func main1() {
	println(banner)
	proxy := &http.Server{
		Addr:    fmt.Sprintf(":%d", Cfg.Port),
		Handler: http.HandlerFunc(handleTunnel),
	}

	log.Println("starting server on port 8765")
	if err := proxy.ListenAndServe(); err != nil {
		log.Fatalf("error starting server: %s", err)
	}
}

func existHost(host string) bool {
	for _, route := range Cfg.Routes {
		if route == host {
			return true
		}
	}
	return false
}

func handleTunnel(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		handleHTTPS(w, r)
	} else {
		handleHTTP(w, r)
	}
}

func handleHTTP(w http.ResponseWriter, r *http.Request) {
	var (
		destConn net.Conn
		//buf        []byte
		targetAddr string
		dialer     *net.Dialer
		//num        int
		err error
	)
	log.Printf("http proxy address: %v\n", r.Host)

	targetAddr = r.Host

	if !existHost(targetAddr) {
		http.Error(w, "http host address not supported", http.StatusServiceUnavailable)
		return
	}

	// 确保地址有端口号
	if !strings.Contains(targetAddr, ":") {
		targetAddr += ":80"
	}

	dialer = &net.Dialer{
		Timeout: 10 * time.Second,
	}
	if destConn, err = dialer.Dial("tcp", targetAddr); err != nil {
		log.Printf("unable to connect to the target server: %v", err)
		http.Error(w, "unable to connect to the target server", http.StatusInternalServerError)
		return
	}
	_ = destConn.SetReadDeadline(time.Now().Add(Cfg.ReadTimeout * time.Second))
	_ = destConn.SetWriteDeadline(time.Now().Add(Cfg.WriteTimeout * time.Second))
	defer destConn.Close()

	if err = r.Write(destConn); err != nil {
		log.Printf("request sending failed: %v", err)
		http.Error(w, "request sending failed", http.StatusInternalServerError)
		return
	}

	if _, err = io.Copy(w, destConn); err != nil {
		log.Printf("request sending failed: %v", err)
	}
}

func handleHTTPS(w http.ResponseWriter, r *http.Request) {

	target := r.Host
	log.Printf("https proxy address: %v\n", r.Host)
	if !strings.Contains(target, ":") {
		target += ":443"
	}

	if !existHost(target) {
		http.Error(w, "https host address not supported", http.StatusServiceUnavailable)
		return
	}

	// 劫持客户端连接
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijacking not supported", http.StatusInternalServerError)
		return
	}
	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer clientConn.Close()

	// 先给客户端返回 200 Connection Established
	if _, err := clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
		return
	}

	// 连目标服务器
	dialer := &net.Dialer{Timeout: 10 * time.Second}
	targetConn, err := dialer.Dial("tcp", target)
	if err != nil {
		return
	}
	defer targetConn.Close()

	// 双向盲转发
	go func() {
		_, _ = io.Copy(targetConn, clientConn)
	}()
	_, _ = io.Copy(clientConn, targetConn)
}
